from torch.utils.data import Dataset
import numpy as np


class DiffusionDataset(Dataset):
    def __init__(self, sets, alpha_bar, step_num=1000):
        super().__init__()
        self.sets = sets
        self.T = step_num
        self.alpha = alpha_bar

    def __len__(self):
        return len(self.sets)
    

    def __getitem__(self, index):
        img, _ = self.sets[index]
        t = np.random.randint(0, self.T)
        e = np.random.normal(0, 1, size=img.shape).astype(np.float32)

        img_weight = self.alpha[t] ** (1 / 2)
        e_weight = (1 - self.alpha[t]) ** (1 / 2)
        x = img_weight * img + e_weight * e

        return x.astype(np.float32), np.array([t + 1], dtype=np.float32), e.flatten()